import matplotlib

matplotlib.use('Agg')
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS']
matplotlib.rcParams['axes.unicode_minus'] = False
from flask import Blueprint, jsonify
from app.models.analysis import AnalysisTask
import numpy as np
import base64
import io
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.decomposition import PCA
from sklearn.metrics import silhouette_samples, silhouette_score

bp = Blueprint('viz', __name__)

# 聚类类别名称映射
CLUSTER_NAMES = {
    0: "高评分新片",
    1: "经典老片",
    2: "小众佳作",
    3: "热门商业片",
    4: "待发掘佳作"
}


@bp.route('/kmeans/<int:task_id>', methods=['GET'])
def kmeans_visualization_api(task_id):
    """根据分析任务ID返回K-means聚类多图可视化结果，直接用analysis结果"""
    task = AnalysisTask.query.get(task_id)
    if not task or not task.result:
        return jsonify({'error': '分析任务不存在或未完成'}), 404
    result = task.result
    labels = np.array(result['labels'])
    centroids = np.array(result['centroids']) if 'centroids' in result else None
    # 如果有PCA降维结果直接用，否则自动降维
    if 'pca' in result:
        reduced = np.array(result['pca'])
    else:
        # 自动降维
        from app.models.dataset import Dataset
        dataset = Dataset.query.get(task.dataset_id)
        data = dataset.get_data()
        pca = PCA(n_components=2)
        reduced = pca.fit_transform(data)
        if centroids is not None:
            centroids = pca.transform(centroids)
    n_clusters = len(np.unique(labels))

    # 散点图
    plt.figure(figsize=(12, 8))
    for i in range(n_clusters):
        plt.scatter(reduced[labels == i, 0], reduced[labels == i, 1], label=CLUSTER_NAMES.get(i, f'聚类 {i}'))
    # 不再绘制聚类中心
    plt.title('K-means聚类可视化')
    plt.xlabel('PC1')
    plt.ylabel('PC2')
    plt.legend()
    scatter_buffer = io.BytesIO()
    plt.savefig(scatter_buffer, format='png', bbox_inches='tight')
    scatter_buffer.seek(0)
    scatter_base64 = base64.b64encode(scatter_buffer.getvalue()).decode()
    plt.close()

    # 轮廓图（如有silhouette，否则动态生成）
    silhouette_base64 = None
    if 'silhouette' in result:
        silhouette_img = base64.b64decode(result['silhouette'])
        silhouette_base64 = result['silhouette']
    else:
        # 动态生成轮廓图
        from sklearn.metrics import silhouette_samples, silhouette_score
        plt.figure(figsize=(12, 8))
        silhouette_avg = silhouette_score(reduced, labels)
        sample_silhouette_values = silhouette_samples(reduced, labels)
        y_lower = 10
        for i in range(n_clusters):
            ith_cluster_silhouette_values = sample_silhouette_values[labels == i]
            ith_cluster_silhouette_values.sort()
            size_cluster_i = ith_cluster_silhouette_values.shape[0]
            y_upper = y_lower + size_cluster_i
            plt.fill_betweenx(np.arange(y_lower, y_upper), 0, ith_cluster_silhouette_values, alpha=0.7)
            plt.text(-0.05, y_lower + 0.5 * size_cluster_i, CLUSTER_NAMES.get(i, f'聚类 {i}'))
            y_lower = y_upper + 10
        plt.axvline(x=silhouette_avg, color="red", linestyle="--")
        plt.title('聚类轮廓分析图')
        plt.xlabel('轮廓系数')
        plt.ylabel('样本')
        silhouette_buffer = io.BytesIO()
        plt.savefig(silhouette_buffer, format='png', bbox_inches='tight')
        silhouette_buffer.seek(0)
        silhouette_base64 = base64.b64encode(silhouette_buffer.getvalue()).decode()
        plt.close()

    # 热力图和雷达图可选（如有中心）
    heatmap_base64 = None
    radar_base64 = None
    if centroids is not None:
        import pandas as pd
        numeric_cols = result.get('feature_names') or [f'特征{i + 1}' for i in range(centroids.shape[1])]
        cluster_centers = pd.DataFrame(centroids, columns=numeric_cols)
        cluster_centers.index = [CLUSTER_NAMES.get(i, f'聚类 {i}') for i in range(n_clusters)]
        # 热力图
        plt.figure(figsize=(12, 8))
        sns.heatmap(cluster_centers, annot=True, cmap='YlOrRd', fmt='.2f')
        plt.title('聚类中心特征热力图')
        heatmap_buffer = io.BytesIO()
        plt.savefig(heatmap_buffer, format='png', bbox_inches='tight')
        heatmap_buffer.seek(0)
        heatmap_base64 = base64.b64encode(heatmap_buffer.getvalue()).decode()
        plt.close()
        # 雷达图
        plt.figure(figsize=(12, 8))
        important_features = numeric_cols[:5]
        angles = np.linspace(0, 2 * np.pi, len(important_features), endpoint=False)
        angles = np.concatenate((angles, [angles[0]]))
        ax = plt.subplot(111, polar=True)
        for i in range(n_clusters):
            values = cluster_centers.iloc[i][important_features].values
            values = np.concatenate((values, [values[0]]))
            ax.plot(angles, values, 'o-', linewidth=2, label=CLUSTER_NAMES.get(i, f'聚类 {i}'))
            ax.fill(angles, values, alpha=0.25)
        plt.xticks(angles[:-1], important_features)
        plt.title('聚类特征雷达图')
        plt.legend(loc='upper right', bbox_to_anchor=(0.1, 0.1))
        radar_buffer = io.BytesIO()
        plt.savefig(radar_buffer, format='png', bbox_inches='tight')
        radar_buffer.seek(0)
        radar_base64 = base64.b64encode(radar_buffer.getvalue()).decode()
        plt.close()
    return jsonify({
        'status': 'success',
        'data': {
            'visualizations': {
                'scatter': scatter_base64,
                'silhouette': silhouette_base64,
                'heatmap': heatmap_base64,
                'radar': radar_base64
            }
        }
    })


@bp.route('/linear_regression/<int:task_id>', methods=['GET'])
def linear_regression_visualization_api(task_id):
    """根据分析任务ID返回线性回归可视化结果"""
    task = AnalysisTask.query.get(task_id)
    if not task or not task.result:
        return jsonify({'error': '分析任务不存在或未完成'}), 404
    result = task.result
    y_true = np.array(result.get('y_true') or result.get('actual_values'))
    y_pred = np.array(result.get('y_pred') or result.get('predictions'))
    if y_true is None or y_pred is None:
        return jsonify({'error': '分析结果缺少必要字段'}), 400
    # 散点图
    plt.figure(figsize=(8, 6))
    plt.scatter(y_true, y_pred, alpha=0.5)
    plt.plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], 'r--', lw=2)
    plt.xlabel('实际值')
    plt.ylabel('预测值')
    plt.title('线性回归：实际值 vs 预测值')
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight')
    buf.seek(0)
    scatter_base64 = base64.b64encode(buf.read()).decode()
    plt.close()
    # 残差图
    plt.figure(figsize=(8, 6))
    residuals = y_true - y_pred
    plt.scatter(y_pred, residuals, alpha=0.5)
    plt.axhline(0, color='red', linestyle='--')
    plt.xlabel('预测值')
    plt.ylabel('残差')
    plt.title('线性回归：残差图')
    buf = io.BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight')
    buf.seek(0)
    residual_base64 = base64.b64encode(buf.read()).decode()
    plt.close()
    return jsonify({
        'status': 'success',
        'data': {
            'visualizations': {
                'scatter': scatter_base64,
                'residual': residual_base64
            }
        }
    })
